package rescuecore2.standard.score;

import java.util.HashSet;
import java.util.Set;

import rescuecore2.Timestep;
import rescuecore2.config.Config;
import rescuecore2.score.AbstractScoreFunction;
import rescuecore2.standard.entities.AmbulanceTeam;
import rescuecore2.standard.entities.BirdMan;
import rescuecore2.standard.entities.Civilian;
import rescuecore2.standard.entities.FireBrigade;
import rescuecore2.standard.entities.PoliceForce;
import rescuecore2.worldmodel.ChangeSet;
import rescuecore2.worldmodel.Entity;
import rescuecore2.worldmodel.EntityID;
import rescuecore2.worldmodel.WorldModel;

/**
 * Score function that measures how quickly civilians are discovered by agents.
 */
public class DiscoveryScoreFunction extends AbstractScoreFunction {
	private Set<EntityID> found;

	/**
	 * Construct a DiscoveryScoreFunction.
	 */
	public DiscoveryScoreFunction() {
		super("Civilian discovery time");
	}

	@Override
	public void initialise(WorldModel<? extends Entity> world, Config config) {
		found = new HashSet<EntityID>();
	}

	@Override
	public double score(WorldModel<? extends Entity> world, Timestep timestep) {
		double sum = 0;
		// Look for agents that observed a civilian
		for (EntityID next : timestep.getAgentsWithUpdates()) {
			Entity agent = world.getEntity(next);
			// Only platoon agents can discover civilians
			if (!isPlatoonAgent(agent)) {
				continue;
			}
			ChangeSet perception = timestep.getAgentPerception(next);
			for (EntityID observedID : perception.getChangedEntities()) {
				// Is it already seen?
				if (found.contains(observedID)) {
					continue;
				}
				Entity e = world.getEntity(observedID);
				if (e instanceof Civilian
						&& !perception.getChangedProperties(observedID)
								.isEmpty()) {
					// Seen a new civilian with at least one updated property.
					found.add(observedID);
					sum += timestep.getTime();
				}
			}
		}
		return sum;
	}

	private boolean isPlatoonAgent(Entity e) {
		return e instanceof FireBrigade || e instanceof BirdMan
				|| e instanceof PoliceForce || e instanceof AmbulanceTeam;
	}
}
